package org.arabidopsis.ahocorasick;

import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Iterator;


public class AhoCorasick {
    private final State root = new State(0);

    private boolean prepared = false;

    public void add(byte[] keyword, Object output) {
        if (this.prepared)
            throw new IllegalStateException(
                    "can't add keywords after prepare() is called");
        State lastState = this.root.extendAll(keyword);
        lastState.addOutput(output);
    }

    public void prepare() {
        prepareFailTransitions();
        this.prepared = true;
    }

    public Iterator search(byte[] bytes) {
        return new Searcher(this, startSearch(bytes));
    }

    private void prepareFailTransitions() {
        Queue q = new Queue();
        for (int i = 0; i < 256; i++) {
            if (this.root.get((byte) i) != null) {
                this.root.get((byte) i).setFail(this.root);
                q.add(this.root.get((byte) i));
            }
        }
        prepareRoot();
        while (!q.isEmpty()) {
            State state = q.pop();
            byte[] keys = state.keys();
            for (int j = 0; j < keys.length; j++) {
                State r = state;
                byte a = keys[j];
                State s = r.get(a);
                q.add(s);
                r = r.getFail();
                while (r.get(a) == null) {
                    r = r.getFail();
                }
                s.setFail(r.get(a));
                s.getOutputs().addAll(r.get(a).getOutputs());
            }
        }
    }

    private void prepareRoot() {
        for (int i = 0; i < 256; i++) {
            if (this.root.get((byte) i) == null)
                this.root.put((byte) i, this.root);
        }
    }

    public boolean forbidden(byte[] bytes, HashMap<String, Integer> keyScore) {
        int score = getScore(new SearchResult(this.root, bytes, 0), keyScore);
        return (score >= 100);
    }

    public String forbiddenStr(byte[] bytes, String t) {
        return replaceStr1(new SearchResult(this.root, bytes, 0), t);
    }

    SearchResult startSearch(byte[] bytes) {
        if (!this.prepared)
            throw new IllegalStateException(
                    "can't start search until prepare()");
        return continueSearch(new SearchResult(this.root, bytes, 0));
    }

    SearchResult continueSearch(SearchResult lastResult) {
        byte[] bytes = lastResult.bytes;
        State state = lastResult.lastMatchedState;
        for (int i = lastResult.lastIndex; i < bytes.length; i++) {
            byte b = bytes[i];
            while (state.get(b) == null)
                state = state.getFail();
            state = state.get(b);
            if (state.getOutputs().size() > 0)
                return new SearchResult(state, bytes, i + 1);
        }
        return null;
    }

    String replaceStr(SearchResult lastResult, String t) {
        byte[] bytes = lastResult.bytes;
        State state = lastResult.lastMatchedState;
        for (int i = lastResult.lastIndex; i < bytes.length; i++) {
            byte b = bytes[i];
            while (state.get(b) == null)
                state = state.getFail();
            state = state.get(b);
            if (state.getOutputs().size() > 0) {
                for (int j = i - state.getDepth() + 1; j < i + 1; j++) {
                    bytes[j] = 42;
                }
            }
        }
        String tmp = new String(bytes);


        return tmp;
    }

    String replaceStr1(SearchResult lastResult, String t) {
        byte[] bytes = lastResult.bytes;
        byte[] tmps = new byte[bytes.length + 100];
        int cur = 0, begin = 0, len = 0;
        State state = lastResult.lastMatchedState;
        for (int i = lastResult.lastIndex; i < bytes.length; i++) {

            byte b = bytes[i];

            while (state.get(b) == null)
                state = state.getFail();
            state = state.get(b);
            if (state.getOutputs().size() > 0) {


                len = i - state.getDepth() - begin + 1;
                if (len > 0) {
                    System.arraycopy(bytes, begin, tmps, cur, len);
                    cur += len;
                }


                System.arraycopy(t.getBytes(), 0, tmps, cur, (t.getBytes()).length);
                cur += (t.getBytes()).length;
                begin = i + 1;
            }
        }
        System.arraycopy(bytes, begin, tmps, cur, bytes.length - begin);
        cur += bytes.length - begin;
        String tmp = new String(tmps, 0, cur, StandardCharsets.UTF_8);
        return tmp;
    }


    int getScore(SearchResult lastResult, HashMap<String, Integer> keyScore) {
        int score = 0;
        byte[] bytes = lastResult.bytes;
        State state = lastResult.lastMatchedState;
        for (int i = lastResult.lastIndex; i < bytes.length; i++) {
            byte b = bytes[i];
            while (state.get(b) == null)
                state = state.getFail();
            state = state.get(b);
            if (state.getOutputs().size() > 0) {
                score += keyScore.get(state.getOutputs().iterator().next()).intValue();
                if (score >= 100 || i + 1 > bytes.length) {
                    break;
                }
            }
        }
        return score;
    }
}